import ast
import json
import javalang
import esprima

VERSION = "v2.3"
NODE_TYPE_DIRECTORY = "directory"
NODE_TYPE_FILE = "file"
NODE_TYPE_CLASS = "class"
NODE_TYPE_FUNCTION = "function"
EDGE_TYPE_CONTAINS = "contains"
EDGE_TYPE_INHERITS = "inherits"
EDGE_TYPE_INVOKES = "invokes"
EDGE_TYPE_IMPORTS = "imports"

VALID_NODE_TYPES = [
    NODE_TYPE_DIRECTORY,
    NODE_TYPE_FILE,
    NODE_TYPE_CLASS,
    NODE_TYPE_FUNCTION,
]
VALID_EDGE_TYPES = [
    EDGE_TYPE_CONTAINS,
    EDGE_TYPE_INHERITS,
    EDGE_TYPE_INVOKES,
    EDGE_TYPE_IMPORTS,
]


class CodeAnalyzer(ast.NodeVisitor):
    def __init__(self, filename):
        self.filename = filename
        self.nodes = []
        self.node_name_stack = []
        self.node_type_stack = []

    def visit_ClassDef(self, node):
        class_name = node.name
        full_class_name = ".".join(self.node_name_stack + [class_name])
        self.nodes.append(
            {
                "name": full_class_name,
                "type": NODE_TYPE_CLASS,
                "code": self._get_source_segment(node),
                "start_line": node.lineno,
                "end_line": node.end_lineno,
            }
        )

        self.node_name_stack.append(class_name)
        self.node_type_stack.append(NODE_TYPE_CLASS)
        self.generic_visit(node)
        self.node_name_stack.pop()
        self.node_type_stack.pop()

    def visit_FunctionDef(self, node):
        if (
            self.node_type_stack
            and self.node_type_stack[-1] == NODE_TYPE_CLASS
            and node.name == "__init__"
        ):
            return
        self._visit_func(node)

    def visit_AsyncFunctionDef(self, node):
        self._visit_func(node)

    def _visit_func(self, node):
        function_name = node.name
        full_function_name = ".".join(self.node_name_stack + [function_name])
        self.nodes.append(
            {
                "name": full_function_name,
                "parent_type": (
                    self.node_type_stack[-1] if self.node_type_stack else None
                ),
                "type": NODE_TYPE_FUNCTION,
                "code": self._get_source_segment(node),
                "start_line": node.lineno,
                "end_line": node.end_lineno,
            }
        )

        self.node_name_stack.append(function_name)
        self.node_type_stack.append(NODE_TYPE_FUNCTION)
        self.generic_visit(node)
        self.node_name_stack.pop()
        self.node_type_stack.pop()

    def _get_source_segment(self, node):
        with open(self.filename, "r") as file:
            source_code = file.read()
        return ast.get_source_segment(source_code, node)


# Parese the given file, use CodeAnalyzer to extract classes and helper functions from the file
def analyze_file(filepath):
    with open(filepath, "r") as file:
        code = file.read()
        # code = handle_edge_cases(code)
        try:
            if filepath.endswith(".py"):
                tree = ast.parse(code, filename=filepath)
            elif filepath.endswith(".java"):
                with open(filepath, "r") as f:
                    content = f.read()
                    tree = javalang.parse.parse(content)
            elif filepath.endswith(".js"):
                tree = esprima.parseScript(filepath)
            else:
                pass
        except:
            raise SyntaxError
    analyzer = CodeAnalyzer(filepath)
    try:
        analyzer.visit(tree)
    except RecursionError:
        pass
    return analyzer.nodes


if __name__ == "__main__":
    file_path = "/data_ext/ref_code/LocAgent/data/data/pp_data/fmxm/run.py"
    # file_path = "/data_ext/ref_code/LocAgent/data/data/pp_data/pingp-auto-test/auto-test-b2b/src/test/java/auto/b2b/check/CheckBusinessAccount.java"
    nodes = analyze_file(file_path)
    print(len(nodes))
    print(json.dumps(nodes[0], ensure_ascii=False, indent=4))
